package cas.ihep.hep.analysis;

import cas.ihep.fs.UnifiedFileSystem;
import cas.ihep.hep.util.Buffer;
import cas.ihep.hep.util.EventBlock;
import cas.ihep.hep.util.MyApplication;
import cas.ihep.hep.util.NativeBlock;
import cas.ihep.spark.io.Packet;
import cas.ihep.spark.io.Serializer;
import cas.ihep.spark.loop.IterationServer;
import cas.ihep.spark.loop.IterativeRunnable;
import cas.ihep.spark.rdd.JIterativePairRDD;
import cas.ihep.spark.util.GlobalConfiguration;
import cas.ihep.spark.util.NamedResourceManager;
import cas.ihep.spark.util.function.Function2;
import cas.ihep.spark.util.function.Function4;
import cas.ihep.util.Finalizer;
import cas.ihep.util.MemoryBuffer;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaSparkContext;
import org.freehep.math.minuit.FCNBase;
import org.freehep.math.minuit.FunctionMinimum;
import org.freehep.math.minuit.MnMigrad;
import org.freehep.math.minuit.MnUserParameters;
import org.luaj.vm2.LuaTable;
import org.luaj.vm2.LuaValue;
import org.luaj.vm2.Varargs;
import scala.Tuple2;
import scala.Tuple3;

import javax.script.*;
import java.io.*;
import java.net.URI;
import java.util.ArrayList;
import java.util.Formatter;
import java.util.List;

public class NewLoopPWA extends MyApplication {
    private MnUserParameters mParameters;
    private Bindings values;

    public NewLoopPWA(String []args) throws IOException, ScriptException {
        super(args);
        mParameters=new MnUserParameters();
        ScriptEngineManager sem=new ScriptEngineManager();
        ScriptEngine engine=sem.getEngineByName("luaj");
        values=new SimpleBindings();
        try(Reader reader=new BufferedReader(new FileReader(args[0]))){
            engine.eval(reader,values);
        }
        LuaTable optionObj=(LuaTable) values.get("option");
        for(LuaValue i = LuaValue.NIL;;){
            Varargs kv=optionObj.next(i);
            i=kv.arg1();
            if(i.isnil()){
                break;
            }
            LuaValue vvv=kv.arg(2);
            if(i.isstring() && vvv.istable()){
                String pname=i.tojstring();
                LuaTable tb=(LuaTable)vvv;
                double v=tb.get(1).todouble();
                double e=tb.get(2).todouble();
                double l=tb.get(3).todouble();
                double u=tb.get(4).todouble();
                mParameters.add(pname,v,e);
                if (e<0){
                    mParameters.fix(pname);
                }else{
                    if(l!=999 && u!=999) {
                        mParameters.setLimits(pname, l, u);
                    }else if(l!=999){
                        mParameters.setLowerLimit(pname,l);
                    }else if(u!=999){
                        mParameters.setUpperLimit(pname,u);
                    }
                }
            }
        }
    }

    @Override
    public void run() throws Exception {
        srun(values,mParameters);
    }

    private static Object bindingsGetOrDefault(Bindings values,String key,Object defVal){
        if(values.containsKey(key)){
            return values.get(key);
        }
        return defVal;
    }

    private static void srun(Bindings values,MnUserParameters mParameters)throws Exception{
        SparkConf sconf=new SparkConf();
        sconf.setAppName("NewLoopPWA");
        LuaTable dataf=(LuaTable) values.get("data");
        LuaTable mcf=(LuaTable)values.get("mc");
        String outputRecord=(String)values.get("record");
        final int partitions=(Integer)values.get("slices");
        GlobalConfiguration.setProperty("spark.loop.type",(String)bindingsGetOrDefault(values,"LoopType","udp"));
        GlobalConfiguration.setProperty("spark.loop.log",(String)bindingsGetOrDefault(values,"LogLevel","Info"));
        try(Finalizer closer= Finalizer.create()){
            PrintWriter output=closer.register(new PrintWriter(new FileWriter(outputRecord)));
            Tuple3<List<Tuple2<EventBlock,EventBlock>>,Long,Long> source=parseInputData(new URI(dataf.get(LuaValue.valueOf("file")).tojstring()).normalize(),
                    new URI(mcf.get(LuaValue.valueOf("file")).tojstring()).normalize(),partitions);
            final long dtNumber=source._2(),mcNumber=source._3();
            JavaSparkContext jsc=closer.register(new JavaSparkContext(sconf));
            JIterativePairRDD<EventBlock,EventBlock> sourcerdd= JIterativePairRDD.from(jsc.parallelize(source._1(),50));
            long beforeStart=System.nanoTime();
            //noinspection Convert2Lambda
            IterativeRunnable<double[],LoopResult> iRun=new IterativeRunnable<double[], LoopResult>() {
                public void run(IterationServer<double[], LoopResult> loopServer) {
                    final int[] accumulator=new int[1];
                    final double[] times=new double[]{(System.nanoTime()-beforeStart)/1e6,0};
                    long start=System.nanoTime();
                    //noinspection Convert2Lambda
                    MnMigrad migrad = new MnMigrad(new FCNBase() {
                        public double valueOf(double[] par) {
                            accumulator[0]++;
                            long start0=System.nanoTime();
                            LoopResult value = loopServer.broadcast(par, new LoopResult(0, 0,0),
                                    addFunction);
                            start0=System.nanoTime()-start0;
                            times[0]+=start0/1e6;
                            times[1]+=value.time/50.0;
                            double ret=value.dt + Math.log(value.mc / mcNumber) * dtNumber;
                            output.println(ret);
                            return ret;
                        }
                    }, mParameters, 1);
                    migrad.setPrecision(1e-11);
                    FunctionMinimum minimum = migrad.minimize(100000, 200);
                    start=System.nanoTime()-start;
                    MnUserParameters resultpars = minimum.userParameters();
                    double pvec2 = resultpars.value(2);
                    double pvec3 = resultpars.value(3);
                    double pvec6 = resultpars.value(6);
                    double pvec7 = resultpars.value(7);
                    new Formatter(System.out).format("Exiting spark application ...\nIsValid = %s\n" +
                                    "likelihood = %f\nUser Parameters: %f %f %f %f\n" +
                                    "Using time: %f ms\nNumber of function calls: %d\nAverage valid computing time=%.02f ms\n"+
                            "Average total time=%.02f ms\nAverage valid computing rate=%.02f\n",
                            minimum.isValid()?"true":"false",minimum.fval(),pvec2,pvec3,pvec6,pvec7,(double)start/1e6,accumulator[0],times[1]/accumulator[0],
                            times[0]/accumulator[0],times[1]/times[0]);
                }
            };
            //noinspection Convert2Lambda
            Function4<EventBlock,EventBlock,double[], NamedResourceManager,LoopResult> clientFunc=
                    new Function4<EventBlock, EventBlock, double[], NamedResourceManager, LoopResult>() {
                        @Override
                        public LoopResult call(EventBlock blk0, EventBlock blk1, double[] param, NamedResourceManager closer) {
                            try {
                                //NewUnsafe.loadLibrary("/mnt/nfs/libNewUnsafe.so");
                                long start=System.nanoTime();
                                NativeBlock nblk0 = (NativeBlock) closer.get("block0");
                                if (nblk0 == null) {
                                    nblk0 = blk0.nativeData();
                                    closer.put("block0",nblk0);
                                }
                                NativeBlock nblk1=(NativeBlock) closer.get("block1");
                                if(nblk1==null){
                                    nblk1=blk1.nativeData();
                                    closer.put("block1",nblk1);
                                }
                                /*LLVMInvoker invoker=(LLVMInvoker)closer.get("llvm");
                                if(invoker==null){
                                    invoker=new LLVMInvoker("/mnt/nfs/newpwa.bc");
                                    closer.put("llvm",invoker);
                                }*/
                                //double[] ret = NewUnsafe.compute(param, nblk0.address(), blk0.count(), nblk1.address(), blk1.count());
                                MemoryBuffer mbuf0=nblk0.getBuffer().mark(),mbuf1=nblk1.getBuffer().mark();
                                double[] ret = NewPWACalculator.compute(param, mbuf0, blk0.count(), mbuf1, blk1.count(),0);
                                /*double[] ret=new double[2];
                                invoker.invoke("NewPwaCompute",param,nblk0.address(),blk0.count(),
                                        nblk1.address(),blk1.count(),ret);*/
                                return new LoopResult(ret[0],ret[1],(System.nanoTime()-start)/1e6);
                            }catch (IOException e){
                                throw new RuntimeException(e);
                            }
                        }
                    };
            sourcerdd.iterate(iRun, clientFunc,double[].class,new LoopResult(0,0,0),addFunction);
            for(Tuple2<EventBlock,EventBlock> t:source._1()){
                t._1().close();
                t._2().close();
            }
        }
    }

    static Tuple3<List<Tuple2<EventBlock,EventBlock>>,Long,Long> parseInput(URI dtfile,URI mcfile,int partitions)throws IOException{
        try(Finalizer closer=Finalizer.create()){
            InputStream dtis=closer.register(UfsHolder.ufs.openStream(dtfile));
            InputStream mcis=closer.register(UfsHolder.ufs.openStream(mcfile));
            byte[] countBuf=new byte[8];
            int result=dtis.read(countBuf);
            if(result<8){
                throw new EOFException("Unexpected EOF");
            }
            long dtcount= Buffer.byteArray2Long(countBuf,0);
            result=mcis.read(countBuf);
            if(result<8){
                throw new EOFException("Unexpected EOF");
            }
            long mccount=Buffer.byteArray2Long(countBuf,0);
            List<EventBlock> dtblks=EventBlock.splitBlock(dtis,8*8,(int)dtcount,partitions);
            List<EventBlock> mcblks=EventBlock.splitBlock(mcis,8*8,(int)mccount,partitions);
            List<Tuple2<EventBlock,EventBlock>> ret=new ArrayList<>(dtblks.size());
            for(int i=0;i<dtblks.size();i++){
                ret.add(new Tuple2<>(dtblks.get(i),mcblks.get(i)));
            }
            return new Tuple3<>(ret,dtcount,mccount);
        }
    }

    static Tuple3<List<Tuple2<EventBlock,EventBlock>>,Long,Long> parseInputData(URI dtfile,URI mcfile,int partitions)throws IOException{
        try(Finalizer closer=Finalizer.create()){
            InputStream dtis=closer.register(UfsHolder.ufs.openStream(dtfile));
            InputStream mcis=closer.register(UfsHolder.ufs.openStream(mcfile));
            byte[] countBuf=new byte[8];
            int result=dtis.read(countBuf);
            if(result<8){
                throw new EOFException("Unexpected EOF");
            }
            long dtcount= Buffer.byteArray2Long(countBuf,0);
            result=mcis.read(countBuf);
            if(result<8){
                throw new EOFException("Unexpected EOF");
            }
            long mccount=Buffer.byteArray2Long(countBuf,0);
            int dteach=(int)(dtcount/partitions);
            if(dtcount%partitions!=0){
                dteach++;
            }
            int mceach=(int)(mccount/partitions);
            if(mccount%partitions!=0){
                mceach++;
            }
            long dttotal=0,mctotal=0;
            ArrayList<Tuple2<EventBlock,EventBlock>> ret=new ArrayList<>();
            for(int i=0;i<partitions;i++){
                int each=Math.min(dteach,(int)(dtcount-dttotal));
                EventBlock b1=EventBlock.readFromFile(UfsHolder.ufs,dtis,each*8*8,each);
                dttotal+=each;
                each=Math.min(mceach,(int)(mccount-mctotal));
                EventBlock b2=EventBlock.readFromFile(UfsHolder.ufs,mcis,each*8*8,each);
                mctotal+=each;
                ret.add(new Tuple2<>(b1,b2));
            }
            return new Tuple3<>(ret,dtcount,mccount);
        }
    }

    private static class UfsHolder{
        private static final UnifiedFileSystem ufs=UnifiedFileSystem.get();
    }

    public static UnifiedFileSystem getUnifiedFileSystem(){
        return UfsHolder.ufs;
    }

    static class LoopParameter implements Serializable{
        double []val;
        long time;
        LoopParameter(double []v){
            val=v;
        }
        static {
            Serializer.registerHandler(LoopParameter.class, new Packet.Serializer() {
                @Override
                public int length(Object o) {
                    LoopParameter lp=(LoopParameter)o;
                    return 4+lp.val.length*8+8;
                }

                @Override
                public void serialize(Object o, DataOutput dataOutput) {
                    try{
                        LoopParameter lp=(LoopParameter)o;
                        dataOutput.writeInt(lp.val.length);
                        for(int i=0;i<lp.val.length;i++){
                            dataOutput.writeDouble(lp.val[i]);
                        }
                        dataOutput.writeLong(lp.time);
                    }catch (IOException ignored){
                        throw new RuntimeException("Should never happens");
                    }
                }

                @Override
                public Object deserialize(DataInput dataInput) {
                    try{
                        int len=dataInput.readInt();
                        double[] v=new double[len];
                        for(int i=0;i<len;i++){
                            v[i]=dataInput.readDouble();
                        }
                        LoopParameter lp=new LoopParameter(v);
                        lp.time=dataInput.readLong();
                        return lp;
                    }catch (IOException ignored){
                        throw new RuntimeException("Should never happens");
                    }
                }
            });
        }
    }

    static class LoopResult implements Serializable {
        double dt,mc;
        double time;
        LoopResult(double v1, double v2,double t){
            dt=v1;
            mc=v2;
            time=t;
        }
        static{
            Serializer.registerHandler(LoopResult.class, new Packet.Serializer() {
                @Override
                public int length(Object data) {
                    return 24;
                }

                @Override
                public void serialize(Object data, DataOutput output) {
                    try{
                        LoopResult result=(LoopResult)data;
                        output.writeDouble(result.dt);
                        output.writeDouble(result.mc);
                        output.writeDouble(result.time);
                    }catch (IOException e){
                        e.printStackTrace();
                    }
                }

                @Override
                public Object deserialize(DataInput input) {
                    try{
                        double v1=input.readDouble();
                        double v2=input.readDouble();
                        double t=input.readDouble();
                        return new LoopResult(v1,v2,t);
                    }catch (IOException e){
                        e.printStackTrace();
                    }
                    return null;
                }
            });
        }

        public String toString(){
            StringBuilder builder=new StringBuilder();
            new Formatter(builder).format("Compute Result<dt: %f, mc: %f>",dt,mc);
            return builder.toString();
        }
    }

    static class LoopResultAdd implements Function2<LoopResult, LoopResult, LoopResult> {
        @Override
        public LoopResult call(LoopResult v1, LoopResult v2) {
            return new LoopResult(v1.dt+v2.dt,v1.mc+v2.mc,v1.time+v2.time);
        }
    }
    private static final LoopResultAdd addFunction=new LoopResultAdd();

}
